from __future__ import division

import copy
import json
import math
import sys
import random

import numpy as np

import utils

class SetEncoder(json.JSONEncoder):
  def default(self, obj):
    if isinstance(obj, set):
      return list(obj)
    return json.JSONEncoder.default(self, obj)

class State:
  global_states_counter = 0
  def __init__(self):
    self.name = 's' + str(State.global_states_counter)
    State.global_states_counter += 1
    self._transitions = {} # map of terminal to list of states (no dups)
    self._transition_count = 0
    self._bits_required = 0
  
  def __repr__(self):
    return self.name

  def add_transition(self, terminal_out, destination_state):
    """ Returns True on success, False if transition already exists. """
    if terminal_out not in self._transitions:
      self._transitions[terminal_out] = []
    if destination_state in self._transitions[terminal_out]:
      return False
    else:
      self._transitions[terminal_out].append(destination_state)
      self._transition_count += 1
      self._bits_required = utils.bits_required(self._transition_count)
      return True
  
  def remove_transition(self, terminal_out, destination_state):
    if terminal_out not in self._transitions:
      return False
    if destination_state not in self._transitions[terminal_out]:
      return False
    self._transitions[terminal_out].remove(destination_state)
    self._transition_count -= 1
    self._bits_required = utils.bits_required(self._transition_count)
    return True    
  
  def get_transitions(self):
    transitions = []
    for k, v in self._transitions.items():
      for state_to in v:
        transitions.append((self, k, state_to))
    return transitions
  
  def get_transition_count(self):
    return self._transition_count

  def __or__(self, other):
    state = State()
    state._transition_count = self._transition_count + other._transition_count
    state._bits_required = utils.bits_required(state._transition_count)
    for k in set(list(other._transitions.keys()) + list(self._transitions.keys())):
      state._transitions[k] = []
      if k in other._transitions:
        state._transitions[k] += other._transitions[k]
      if k in self._transitions:
        state._transitions[k] += self._transitions[k]
      state._transitions[k] = list(set(state._transitions[k]))
    return state

class NFA:
  def __init__(self, terminals, initial_state=None):
    if not initial_state:
      initial_state = State()
    self._initial = initial_state
    self.terminals = set(terminals)
    self.states = utils.IndexedSet([initial_state])
    self._accepting = set()

  def __or__(self, other):
    if self is other:
      return copy.deepcopy(self)
    nfa = NFA(self.terminals | other.terminals, self.get_initial() | other.get_initial())
    nfa.states |= self.states | other.states
    nfa._accepting |= self._accepting | other._accepting
    if self.get_initial() in self._accepting or other.get_initial() in other._accepting:
      nfa._accepting.add(nfa.get_initial())
    nfa.remove_islands()
    return nfa

  def get_initial(self):
    return self._initial

  def description_length(self):
    bits_for_one_state = utils.bits_required(len(self.states))
    bits_for_one_terminal = utils.bits_required(len(self.terminals))
    dl = 0
    dl += len(self.states) + 1 # declare amount of states with 1s, end with a single 0
    dl += (len(self._accepting) + 1) * bits_for_one_state # declare section amount of accepting states and their indices
    dl += (2 * bits_for_one_state + bits_for_one_terminal) * sum([x.get_transition_count() for x in self.states]) # transitions, each transition is a triplet: two states and a terminal
    return dl

  def get_transitions(self):
    transitions = []
    for state in self.states:
      transitions += state.get_transitions()
    return transitions

  def add_state(self):
    new_state = State()
    self.states.add(new_state)
    return new_state

  def add_accepting(self, state):
    assert(state in self.states)
    self._accepting.add(state)

  def get_accepting_indices(self):
    return [self.states.index(x) for x in self._accepting]

  def add_random_state(self):
    state_from = random.choice(self.states)
    state_to = random.choice(self.states)
    new_state = self.add_state()
    terminal_in = utils.set_choice(self.terminals)
    terminal_out = utils.set_choice(self.terminals)
    state_from.add_transition(terminal_in, new_state)
    new_state.add_transition(terminal_out, state_to)
    return 'added state {}, with {} in from {} and {} out to {}'.format(new_state, terminal_in, state_from, terminal_out, state_to)
  def remove_random_state(self):
    if len(self.states) < 2:
      return False
    choice = random.choice(self.states - utils.IndexedSet([self.get_initial()]))
    self.remove_state(choice)
    return 'removed state {}'.format(choice)
  def add_random_transition(self):
    state_from = random.choice(self.states)
    terminal = utils.set_choice(self.terminals)
    state_to = random.choice(self.states)
    result = state_from.add_transition(terminal, state_to)
    if result:
      return 'added transition {} - {} -> {}'.format(state_from, terminal, state_to)
    else:
      return False
  def remove_random_transition(self):
    if len(self.states) < 2:
      return False
    state_from, terminal, state_to = random.choice(self.get_transitions())
    assert(state_from.remove_transition(terminal, state_to))
    return 'removed transition {} - {} -> {}'.format(state_from, terminal, state_to)
  def replace_random_terminal(self):
    state_from, terminal, state_to = random.choice(self.get_transitions())
    assert(state_from.remove_transition(terminal, state_to))
    new_terminal = random.choice([x for x in self.terminals if x != terminal])
    result = state_from.add_transition(new_terminal, state_to)
    if result:
      return 'replaced transition {} - {} -> {} to {} - {} -> {}'.format(state_from, terminal, state_to, state_from, new_terminal, state_to)
    else:
      return False
  def add_random_accepting(self):
    state = random.choice(self.states)
    if state in self._accepting:
      return False
    else:
      self._accepting.add(state)
      return 'added {} to accepting states'.format(state)
  def remove_random_accepting(self):
    state = utils.set_choice(self._accepting)
    if state is self.get_initial():
      return False
    else:
      self._accepting.remove(state)
      return 'removed {} from accepting states'.format(state)

  def remove_state(self, state):
    assert(state is not self.get_initial())
    self._accepting.discard(state)
    self.states.remove(state)
    for s in self.states:
      for t in self.terminals:
        s.remove_transition(t, state)

  def remove_islands(self):
    """ Removes unreachable states """
    states_appearing_in_transitions = utils.IndexedSet()
    for _, _, state in self.get_transitions():
      states_appearing_in_transitions.add(state)
    states_to_remove = self.states - states_appearing_in_transitions
    states_to_remove.discard(self.get_initial())
    for s in states_to_remove:
      self.remove_state(s)
    return len(states_to_remove) + (self.remove_islands() if len(states_to_remove) > 0 else 0)

  def get_neighbor(self):
    nfa_copy = copy.deepcopy(self)
    while True:
      transformation = random.choice([nfa_copy.replace_random_terminal,
                                      nfa_copy.add_random_state,
                                      nfa_copy.remove_random_state,
                                      nfa_copy.add_random_transition,
                                      nfa_copy.remove_random_transition,
                                      nfa_copy.add_random_accepting,
                                      nfa_copy.remove_random_accepting])
      message = transformation()
      if message:
        break
    nfa_copy.remove_islands()
    return nfa_copy, message

  def to_json(self):
    result = {
        'terminals': self.terminals,
        'transitions': [
          {
            's': str(t[0]), # source
            't': t[1], # terminal
            'd': str(t[2]) # destination
          }
          for t in self.get_transitions()],
        'initial': str(self.get_initial()),
        'accepting': [str(s) for s in self._accepting]
    }
    return json.dumps(result, cls=SetEncoder)

  def to_graphviz(self):
    output = 'digraph NFA {\n'
    json_dict = json.loads(self.to_json())
    for transition in json_dict['transitions']:
      output += '\t' + transition['s'] + ' -> ' + transition['d'] + ' [label="' + transition['t'] + '"];\n'
    for accepting in json_dict['accepting']:
      output += '\t' + accepting + ' [peripheries=2];\n'
    output += '\t' + json_dict['initial'] + ' [shape=triangle];\n'
    output += '}'
    print(output)


  def parse(self, input_string):
    # with utils.Profiler():
    # import pprint
    # pp = pprint.PrettyPrinter(indent=4)
    n = len(input_string)
    k = len(self.states)
    # a numpy matrix where in cell i, j is minimal description length to reach state j before scanning input_string[i], or -1 if it's unreachable
    L = np.empty((n+1, len(self.states)))
    L.fill(np.inf)
    # F = [None] * (n + 1) # a list where in position i is a set of reachable states before scanning input_string[i]
    # B = [None] * n # Backtracking: a list of n maps, where the map in position i is between target state and set of origin states when scanning input_string[i]
    # L = [None] * (n + 1) # list where in position i is a map between states and minimal description length to reach them before scanning input_string[i]
    L[0, self.states.index(self._initial)] = 0
    # L[0] = {initial: 0}
    for i in range(0, n):
      # B[i] = {}
      # L[i+1] = {}
      for j in range(k):
        if L[i, j] == np.inf:
          # unreachable state
          continue
        state = self.states[j]
        dl = L[i, j] + state._bits_required # bits required to move from this state
        if input_string[i] in state._transitions:
          for next_state in state._transitions[input_string[i]]:
            # F[i+1].add(next_state)
            # B[i].get(next_state, set()).add(state)
            next_index = self.states.index(next_state)
            if dl < L[i+1, next_index]:
              L[i+1, next_index] = dl
            # print(str(next_state) + ' is reachable from ' + str(state) + ' after reading position ' + str(i))
    # pp.pprint(F)
    # pp.pprint(B)
    min_dl = np.min(L[n, self.get_accepting_indices()])
      # print('bps {}'.format(min([L[n][state] for state in accepting_and_reachable_at_end]) / n))
      #return min([L[n][state] for state in accepting_and_reachable_at_end])
    if min_dl != np.inf:
      return min_dl
    else:
      return False

def test_parsing():
  a = NFA(['a', 'b', 'c'])
  s0 = a.get_initial()
  s1 = a.add_state()
  a.add_accepting(s0)
  s0.add_transition('a', s0)
  s0.add_transition('b', s1)
  s1.add_transition('a', s0)
  s1.add_transition('b', s1)
  def tester(word, expected):
    print('checking word: {}'.format(word))
    print('result is: {}, expected is {}'.format(a.parse(word), expected))
  tester('ababab', False)
  tester('abababa', 7)
  tester('abcabcab', False)
  tester('a', 1)
  tester('', 0)
  tester('b', False)
  tester('c', False)
  tester('ab', False)
  tester('ac', False)
  tester('aa', 2)

def create_initial_nfa(terminals):
  nfa = NFA(terminals)
  states = [nfa.get_initial()]
  for t in terminals[1:]:
    new_state = nfa.add_state()
    states.append(new_state)
  print(states)
  for s in states:
    nfa.add_accepting(s)
    for i in range(len(terminals)):
      s.add_transition(terminals[i], states[i])
  return nfa

def from_json(json_str_or_obj):
  if isinstance(json_str_or_obj, str):
    json_obj = json.loads(json_str)
  else:
    json_obj = json_str_or_obj
  nfa = NFA(json_obj['terminals'])
  name_mapping = {}
  name_mapping[json_obj['initial']] = nfa.get_initial()
  names = set([json_obj['initial']])
  for t in json_obj['transitions']:
    names.add(t['s'])
    names.add(t['d'])
  for name in names:
    if name not in name_mapping:
      s = nfa.add_state()
      name_mapping[name] = s
  for t in json_obj['transitions']:
    name_mapping[t['s']].add_transition(t['t'], name_mapping[t['d']])
  for a in json_obj['accepting']:
    if a not in name_mapping:
      s = nfa.add_state()
      name_mapping[a] = s
    nfa._accepting.add(name_mapping[a])
  return nfa